import SimpleITK as sitk
import os
import numpy as np

"""
resample
"""

def resampleVolume(outsize,vol):
    """
    resample volume to different spacing size\n
    paras：
    outpacing：spacing size，such as [1,1,1]
    vol：sitk image\n
    return：resampled volume
    """
    # outsize = [0,0,0]

    #read
    outspacing = vol.GetSpacing()
    inputsize = outsize
    inputspacing = vol.GetSpacing()

    transform = sitk.Transform()
    transform.SetIdentity()
    #calculate out size
    # outsize[0] = int(inputsize[0]*inputspacing[0]/outspacing[0] + 0.5)
    # outsize[1] = int(inputsize[1]*inputspacing[1]/outspacing[1] + 0.5)
    # outsize[2] = int(inputsize[2]*inputspacing[2]/outspacing[2] + 0.5)

    #out parameters
    resampler = sitk.ResampleImageFilter()
    resampler.SetTransform(transform)
    resampler.SetInterpolator(sitk.sitkLinear)
    resampler.SetOutputOrigin(vol.GetOrigin())
    resampler.SetOutputSpacing(outspacing)
    resampler.SetOutputDirection(vol.GetDirection())
    resampler.SetSize(outsize)
    newvol = resampler.Execute(vol)
    return newvol

def resample_image(itk_image, is_label=False):
    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()
    print(original_spacing)
    out_spacing = [int(round(2*o)) for o in original_spacing]
    print(out_spacing)

    # 上述也可以直接用下面这句简写
    out_size = [int(round(osz*ospc/nspc)) for osz,ospc,nspc in zip(original_size, original_spacing, out_spacing)]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    if is_label: # 如果是mask图像，就选择sitkNearestNeighbor这种插值
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else: # 如果是普通图像，就采用sitkBSpline插值法
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

def resampleSize(sitkImage, new_size):
    #重采样函数
    euler3d = sitk.Euler3DTransform()

    xsize, ysize, zsize = sitkImage.GetSize()
    xspacing, yspacing, zspacing = sitkImage.GetSpacing()
    new_xsize, new_ysize, depth = new_size
    new_spacing_z = zspacing/(depth/float(zsize))
    new_spacing_x = xspacing/(new_xsize/float(xsize))
    new_spacing_y = yspacing / (new_ysize / float(ysize))

    origin = sitkImage.GetOrigin()
    direction = sitkImage.GetDirection()
    #根据新的spacing 计算新的size
    newsize = (int(xsize*xspacing/new_spacing_x),int(ysize*yspacing/new_spacing_y),int(zsize*zspacing/new_spacing_z))
    newspace = (new_spacing_x, new_spacing_y, new_spacing_z)
    sitkImage = sitk.Resample(sitkImage,newsize,euler3d,sitk.sitkNearestNeighbor,origin,newspace,direction)
    return sitkImage

def resample_volume(volume_path, interpolator = sitk.sitkLinear, new_spacing = [1, 1, 1]):
    volume = sitk.ReadImage(volume_path)
    original_spacing = volume.GetSpacing()
    original_size = volume.GetSize()
    new_size = [int(round(osz*ospc/nspc)) for osz,ospc,nspc in zip(original_size, original_spacing, new_spacing)]
    resampled_image = sitk.Resample(volume, new_size, sitk.Transform(), interpolator,
                         volume.GetOrigin(), new_spacing, volume.GetDirection(), 0,
                         volume.GetPixelID())
    sitk.WriteImage(resampled_image,"002.nii.gz")

def adjustMethod(data_resampled,w_width,w_center):
    val_min = w_center - (w_width / 2)
    val_max = w_center + (w_width / 2)
    print(val_max,val_min)

    data_adjusted = data_resampled.copy()
    data_adjusted[data_resampled < val_min] = val_min
    data_adjusted[data_resampled > val_max] = val_max

    return data_adjusted


def downsample(input_dir, tgt_dir, temp_info_dir, seq_list, is_train=False):
    
    for case in seq_list:
        print(case)
        vol_path = os.path.join(input_dir, "case_{}/imaging.nii.gz".format(str(case).zfill(5)))  #autodl-tmp/new_traindata/case_00144/imaging.nii.gz
        volume = sitk.ReadImage(vol_path)
        newvol = resampleSize(volume,[256,256,48])
        data_dir = os.path.join(tgt_dir, 'case_{}'.format(str(case).zfill(5)))
        # data_dir = '/root/autodl-tmp/resize_data/case_'+str(case).zfill(5)
        if not os.path.isdir(data_dir):
            os.makedirs(data_dir)
        sitk.WriteImage(newvol, os.path.join(data_dir, 'imaging.nii.gz'))
        if is_train:
            seg_path = os.path.join(input_dir, "case_{}/segmentation.nii.gz".format(str(case).zfill(5)))
            segmentation = sitk.ReadImage(seg_path)
            newseg = resampleSize(segmentation,[256,256,48])
            sitk.WriteImage(newseg, os.path.join(data_dir, 'segmentation.nii.gz'))
       
        
        # store the template info
        with open(os.path.join(temp_info_dir, 'template_{}.txt'.format(str(case).zfill(5))), 'w') as f:
            
            f.write(str(volume.GetSize()) + '\n')
            f.write(str(volume.GetSpacing()) + '\n')
            f.write(str(volume.GetOrigin()) + '\n')
            f.write(str(volume.GetDirection()) + '\n')
        
        

    
#     for case in [186]:
#         print(case)
#         vol_path = "/root/autodl-tmp/predictions/prediction_{}.nii.gz".format(str(case).zfill(5))  #autodl-tmp/new_traindata/case_00144/imaging.nii.gz
#         # seg_path = "/root/autodl-tmp/new_traindata/case_{}/segmentation.nii.gz".format(str(case).zfill(5))
#         volume = sitk.ReadImage(vol_path)
#         # segmentation = sitk.ReadImage(seg_path)
#         imaging = sitk.GetArrayFromImage(sitk.ReadImage('/root/autodl-tmp/all_traindata/case_{}/imaging.nii.gz'.format(str(case).zfill(5))))
#         size = [imaging.shape[1],imaging.shape[2],imaging.shape[0]]
#         print(size)
#         newvol = resampleSize(volume,size)
#         # vol_array = sitk.GetArrayFromImage(newvol)
#         # new_array = adjustMethod(vol_array,w_width,w_center)
#         # new_array = new_array.astype(dtype=np.float32)
#         # newvol = sitk.GetImageFromArray(new_array)
#         # newseg = resampleSize(segmentation,[256,256,48])
        
#         data_dir = '/root/autodl-tmp/predict_data/'
#         if not os.path.isdir(data_dir):
#             os.makedirs(data_dir)
#         sitk.WriteImage(newvol,'/root/autodl-tmp/predict_data/{}_segmentation.nii.gz'.format(str(case).zfill(5)))
#         # sitk.WriteImage(newseg,'/root/autodl-tmp/resize_data/case_{}/segmentation.nii.gz'.format(str(case).zfill(5)))


def recover(source_dir, template_dir, target_dir, seq_list):
    for case in seq_list:
        print(case)
        img_name = 'prediction_{}.nii.gz'.format(str(case).zfill(5))
        src_img = sitk.ReadImage(os.path.join(source_dir, img_name))
        template_img = sitk.ReadImage(os.path.join(template_dir, r'case_{}/imaging.nii.gz'.format(str(case).zfill(5))))
        src_size = src_img.GetSize()
        tpl_size = template_img.GetSize()
        tpl_spacing = template_img.GetSpacing()
        src_spacing = np.asarray(tpl_spacing) * np.asarray(tpl_size) / np.asarray(src_size) * np.asarray([-1, -1, 1])
        src_img.SetSpacing(src_spacing)
        target_img = sitk.Resample(src_img, tpl_size, sitk.Euler3DTransform(), sitk.sitkNearestNeighbor, template_img.GetOrigin(), tpl_spacing, template_img.GetDirection())
        
        sitk.WriteImage(target_img, os.path.join(target_dir, r'prediction_{}.nii.gz'.format(str(case).zfill(5))))
        
    pass



def static_windows(data, windows, floor=0, ceiling=1):
    """
    Normalizes the data according to a set of predefined windows. This is helpful for CT normalization where the
    units are static and radiologists often have a set of windowing parameters that the use that allow them to look at
    different features in the image.
    :param data: 3D numpy array.
    :param windows:
    :param floor: defaults to 0.
    :param ceiling: defaults to 1.
    :return: Array with data windows listed in the final dimension
    """
    data = np.squeeze(data)
    normalized_data = np.ones(data.shape + (len(windows),)) * floor
    for i, (l, w) in enumerate(windows):
        normalized_data[..., i] = radiology_style_windowing(data, l, w, floor=floor, ceiling=ceiling)
    return normalized_data


def radiology_style_windowing(data, l, w, floor=0, ceiling=1):
    upper = l + w/2
    lower = l - w/2
    return window_data(data, floor_threshold=lower, ceiling_threshold=upper, floor=floor, ceiling=ceiling)


def window_data(data, floor_threshold, ceiling_threshold, floor, ceiling):
    data = (data - floor_threshold) / (ceiling_threshold - floor_threshold)
    # set the data below the floor to equal the floor
    data[data < floor] = floor
    # set the data above the ceiling to equal the ceiling
    data[data > ceiling] = ceiling
    return data




def normalize(src_root, windows, seq_list):
    for seq_number in seq_list:
        src_path = os.path.join(src_root, 'case_' + seq_number, 'imaging.nii.gz')
        src_img = sitk.ReadImage(src_path)
        src_array = sitk.GetArrayFromImage(src_img)
        
        normalized_data = static_windows(src_array, windows)
        normalized_data = np.squeeze(normalized_data)
        normalized_img = sitk.GetImageFromArray(normalized_data)
        normalized_img.CopyInformation(src_img)
        
        sitk.WriteImage(normalized_img, src_path)
        print('finished normalize,', seq_number)


def window_normalize(data, window=[(400, 1200)]):
    normalized_data = static_windows(data, window)
    normalized_data = np.squeeze(normalized_data)
    return normalized_data


if __name__ == '__main__':
    # input the seq_list

    train_list = [144, 146, 147, 149, 152, 154, 155, 158, 159, 160, 161, 162, 166, 167, 168, 170, 171, 172, 173, 176]
    train_list = [str(x).zfill(5) for x in train_list]
    train_list += [str(x).zfill(5) for x in range(200, 258)]

    validation_list = [str(x).zfill(5) for x in range(177, 197)]
    validation_list.remove('00186')

    test_list = [str(i).zfill(5) for i in range(260, 345)]

    src_root = '/root/autodl-tmp/all_traindata'
    tgt_root = '/root/autodl-tmp/resize_data_normalized'
    # tgt_root = '/root/autodl-tmp/resize_data'
    template_info_root = '/root/autodl-tmp/longaxis_proj/longaxis_calc/mydata_test/pred_recover/template_info/'
    downsample(src_root, tgt_root, template_info_root, train_list + validation_list, is_train=True)
    normalize(tgt_root, [(400, 1200)], train_list + validation_list)
    
